package com.jl.magic.extract;

import java.util.*;

/**
 * @author jiangli
 */
public class MMCQ {

    /**
     * 像素rgb分量区间化后占位数
     */
    public static final int QUANTIZE_HOLD_BITS = 5;
    /**
     * 像素rgb分量区间化位移的位数
     */
    public static final int QUANTIZE_MOVE_BITS = 8 - QUANTIZE_HOLD_BITS;
    /**
     * 像素rgb分量区间化后，像素立方的体积
     */
    public static final int CUBE_SIZE = 1 << (3 * QUANTIZE_HOLD_BITS);
    /**
     * 像素rgb分量区间化后，像素立方的最大边长
     */
    private static final int CUBE_LENGTH = 1 << QUANTIZE_HOLD_BITS;
    /**
     * 目标取色负载因子
     */
    private static final double FACTOR_OF_TARGET = 0.75;
    /**
     * 最大循环切割次数
     */
    private static final int MAX_LOOP = 1000;

    /**
     * 主题色最多提取个数
     */
    private static final int MAX_EXTRACT_COUNT = 10;

    /**
     * 获取一个缩小后像素的索引
     *
     * @param r red 分量
     * @param g green 分量
     * @param b blue 分量
     * @return 索引
     */
    public static int getColorIndex(int r, int g, int b) {
        return (r << (2 * QUANTIZE_HOLD_BITS)) + (g << QUANTIZE_HOLD_BITS) + b;
    }

    /**
     * 提取主题色
     *
     * @param pixels    像素
     * @param maxColors 主题色个数
     * @return 主题色map
     */
    public static ColorMap quantize(int[][] pixels, int maxColors) {

        if (pixels.length == 0 || maxColors < 1 || maxColors > MAX_EXTRACT_COUNT) {
            return new ColorMap();
        }
        //像素在量子化空间的数量
        int[] histories = getHistories(pixels);

        //构建初始像素立方PixelCube
        PixelCube pixelCube = buildPixelCube(pixels, histories);

        List<PixelCube> cubes = new ArrayList<>();
        cubes.add(pixelCube);

        //目标取色数
        int target = (int) Math.ceil(FACTOR_OF_TARGET * maxColors);

        loopCut(cubes, COMPARATOR_COUNT, target, histories);

        cubes.sort(COMPARATOR_PRODUCT);

        //下一组—使用(npix * vol)排序生成中值切割。
        if (maxColors > cubes.size()) {
            loopCut(cubes, COMPARATOR_PRODUCT, maxColors, histories);
        }

        //反转排序cubes列表，将有最多元素的放前面
        Collections.reverse(cubes);

        //计算实际颜色
        ColorMap colorMap = new ColorMap();
        for (PixelCube cube : cubes) {
            colorMap.push(cube);
        }
        return colorMap;
    }

    /**
     * 循环切分
     *
     * @param cubes      像素立方体列表
     * @param comparator 像素数量比较器
     * @param target     目标取色
     * @param histories  量化后的像素在一维空间的分布
     */
    private static void loopCut(List<PixelCube> cubes, Comparator<PixelCube> comparator, int target, int[] histories) {
        int loop = 0;
        PixelCube pixelCube;
        while (loop < MAX_LOOP) {
            pixelCube = cubes.get(cubes.size() - 1);
            if (pixelCube.count() == 0) {
                cubes.sort(comparator);
                loop++;
                continue;
            }
            cubes.remove(cubes.size() - 1);

            //切分操作
            PixelCube[] pixelCubes = medianCutApply(histories, pixelCube);
            PixelCube cube1 = pixelCubes[0];
            PixelCube cube2 = pixelCubes[1];

            if (cube1 == null) {
                throw new RuntimeException("切分出错! 应该不会发生！");
            }

            cubes.add(cube1);
            if (cube2 != null) {
                cubes.add(cube2);
            }
            cubes.sort(comparator);

            if (cubes.size() >= target) {
                return;
            }
            loop++;
        }
    }

    /**
     * 中值切分
     *
     * @param histories 量化后的像素在一维空间的分布
     * @param pixelCube 待切分的像素立方
     * @return 切分后的像素立方
     */
    private static PixelCube[] medianCutApply(int[] histories, PixelCube pixelCube) {

        //只有一个像素，不再切分
        if (pixelCube.count() == 1) {
            return new PixelCube[]{pixelCube.clone(), null};
        }

        int redWidth = pixelCube.r2 - pixelCube.r1 + 1;
        int greenWidth = pixelCube.g2 - pixelCube.g1 + 1;
        int blueWidth = pixelCube.b2 - pixelCube.b1 + 1;
        //确定最长边，按最长边进行切分
        int maxWidth = Math.max(Math.max(redWidth, greenWidth), blueWidth);

        // 沿着选定的轴查找局部总和数组.
        //总像素数个数
        int total = 0;
        //第i的位置存放为[0,i]面的空间里像素点总和，i >= 0 且 i < partialSum.length
        int[] partialSum = new int[CUBE_LENGTH];
        // -1 = not set / 0 = 0
        Arrays.fill(partialSum, -1);
        //第i的位置存放为(i,CUBE_LENGTH-1]面的空间里像素总和
        int[] lookAheadSum = new int[CUBE_LENGTH];
        // -1 = not set / 0 = 0
        Arrays.fill(lookAheadSum, -1);
        int i, j, k, sum, index;

        if (maxWidth == redWidth) {
            for (i = pixelCube.r1; i <= pixelCube.r2; i++) {
                sum = 0;
                for (j = pixelCube.g1; j <= pixelCube.g2; j++) {
                    for (k = pixelCube.b1; k <= pixelCube.b2; k++) {
                        index = getColorIndex(i, j, k);
                        sum += histories[index];
                    }
                }
                total += sum;
                partialSum[i] = total;
            }
        } else if (maxWidth == greenWidth) {
            for (i = pixelCube.g1; i <= pixelCube.g2; i++) {
                sum = 0;
                for (j = pixelCube.r1; j <= pixelCube.r2; j++) {
                    for (k = pixelCube.b1; k <= pixelCube.b2; k++) {
                        index = getColorIndex(j, i, k);
                        sum += histories[index];
                    }
                }
                total += sum;
                partialSum[i] = total;
            }
        } else {
            /* maxWidth == blueWidth */
            for (i = pixelCube.b1; i <= pixelCube.b2; i++) {
                sum = 0;
                for (j = pixelCube.r1; j <= pixelCube.r2; j++) {
                    for (k = pixelCube.g1; k <= pixelCube.g2; k++) {
                        index = getColorIndex(j, k, i);
                        sum += histories[index];
                    }
                }
                total += sum;
                partialSum[i] = total;
            }
        }

        for (i = 0; i < CUBE_LENGTH; i++) {
            if (partialSum[i] != -1) {
                lookAheadSum[i] = total - partialSum[i];
            }
        }

        //确定切割平面
        return maxWidth == redWidth ? doCut('r', pixelCube, partialSum, lookAheadSum, total)
                : maxWidth == greenWidth ? doCut('g', pixelCube, partialSum, lookAheadSum, total)
                : doCut('b', pixelCube, partialSum, lookAheadSum, total);
    }

    private static PixelCube[] doCut(char color, PixelCube pixelCube, int[] partialSum,
                                     int[] lookAheadSum, int total) {
        int min;
        int max;

        if (color == 'r') {
            min = pixelCube.r1;
            max = pixelCube.r2;
        } else if (color == 'g') {
            min = pixelCube.g1;
            max = pixelCube.g2;
        } else {
            /* color == 'b' */
            min = pixelCube.b1;
            max = pixelCube.b2;
        }

        int left, right;
        PixelCube cube1, cube2;
        int cutPoint, count2;

        for (int i = min; i <= max; i++) {
            if (partialSum[i] > total / 2) {
                cube1 = pixelCube.clone();
                cube2 = pixelCube.clone();

                //左侧长度
                left = i - min;
                //右侧长度
                right = max - i;

                if (left <= right) {
                    cutPoint = Math.min(max - 1, (i + right / 2));
                } else {
                    // 2.0 and cast to int is necessary to have the same behaviour as in JavaScript
                    cutPoint = Math.max(min, ((int) (i - 1 - left / 2.0)));
                }

                // avoid 0-count boxes
                while (cutPoint < 0 || partialSum[cutPoint] <= 0) {
                    cutPoint++;
                }
                count2 = lookAheadSum[cutPoint];
                while (count2 == 0 && cutPoint > 0 && partialSum[cutPoint - 1] > 0) {
                    count2 = lookAheadSum[--cutPoint];
                }

                // set dimensions
                if (color == 'r') {
                    cube1.r2 = cutPoint;
                    cube2.r1 = cutPoint + 1;
                } else if (color == 'g') {
                    cube1.g2 = cutPoint;
                    cube2.g1 = cutPoint + 1;
                } else {
                    /* color == 'b' */
                    cube1.b2 = cutPoint;
                    cube2.b1 = cutPoint + 1;
                }

                return new PixelCube[]{cube1, cube2};
            }
        }

        throw new RuntimeException("像素空间不能够切分!");
    }

    //比较两个像素空间里实际像素的数量
    private static final Comparator<PixelCube> COMPARATOR_COUNT = Comparator.comparingInt(PixelCube::count);

    private static final Comparator<PixelCube> COMPARATOR_PRODUCT = (a, b) -> {
        int aCount = a.count();
        int bCount = b.count();
        int aVolume = a.volume();
        int bVolume = b.volume();
        //如果像素个数相等，则按体量排序
        if (aCount == bCount) {
            return aVolume - bVolume;
        }
        //否则按 count * volume 排序
        return Long.compare((long) aCount * aVolume, (long) bCount * bVolume);
    };

    /**
     * 构建初始的PixelCube
     *
     * @param pixels    像素点
     * @param histories 量化的像素点分布
     * @return 像素立方体对象
     */
    private static PixelCube buildPixelCube(int[][] pixels, int[] histories) {
        int rMin = 1000000, rMax = 0;
        int gMin = 1000000, gMax = 0;
        int bMin = 1000000, bMax = 0;

        for (int[] ints : pixels) {
            int[] pixel = ints.clone();
            int[] rgb = interval(pixel);

            if (rgb[0] < rMin) {
                rMin = rgb[0];
            } else if (rgb[0] > rMax) {
                rMax = rgb[0];
            }

            if (rgb[1] < gMin) {
                gMin = rgb[1];
            } else if (rgb[1] > gMax) {
                gMax = rgb[1];
            }

            if (rgb[2] < bMin) {
                bMin = rgb[2];
            } else if (rgb[2] > bMax) {
                bMax = rgb[2];
            }
        }

        return new PixelCube(rMin, rMax, gMin, gMax, bMin, bMax, histories);
    }

    /**
     * @param pixels 像素点数组
     * @return 一维数组，表示颜色空间中每个量子化区域的像素数
     */
    private static int[] getHistories(int[][] pixels) {
        int[] histories = new int[CUBE_SIZE];
        int index;
        for (int[] ints : pixels) {
            int[] pixel = ints.clone();
            int[] rgb = interval(pixel);
            index = getColorIndex(rgb[0], rgb[1], rgb[2]);
            histories[index]++;
        }
        return histories;
    }

    /**
     * 像素区间化
     *
     * @param rgb 像素点rgb值
     * @return 区间化后的像素rgb值
     */
    private static int[] interval(int[] rgb) {
        for (int i = 0; i < rgb.length; i++) {
            rgb[i] = rgb[i] >> QUANTIZE_MOVE_BITS;
        }
        return rgb;
    }

}
